import time
import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
from torch.utils.data import Dataset
torch.cuda.empty_cache()


#对标签进行独热编码
def convert_to_one_hot(Y, C):
    Y = Y.astype(int)  # 确保Y是整数类型
    #Y = Y.reshape(-1, 1)
    one_hot = np.zeros((Y.shape[0], C), dtype=int)
    for i in range(Y.shape[0]):
        one_hot[i][Y[i]]=1
    return one_hot.T

def model_s_construct():
    class CNN(nn.Module):
        def __init__(self):
            super(CNN,self).__init__()
            #第一个卷积通道
            self.conv11 = nn.Conv1d(4, 16, kernel_size=2, padding=1)
            self.conv12 = nn.Conv1d(16, 32, 3, padding=1)
            self.conv13 = nn.Conv1d(32, 64, 3, padding=1)
            self.pool11 = nn.MaxPool1d(2,2)
            self.pool12 = nn.MaxPool1d(2,2)

            #第二个卷积通道
            self.conv21 = nn.Conv1d(4, 32, 2, padding=2)
            self.conv22 = nn.Conv1d(32, 64, 2, padding=1)
            self.conv23 = nn.Conv1d(64, 128, 3, padding=1)
            self.pool21 = nn.MaxPool1d(2,2)
            self.pool22 = nn.MaxPool1d(3,3)

            #第三个卷积通道
            self.conv31 = nn.Conv1d(4, 16, 2, padding=1)
            self.conv32 = nn.Conv1d(16, 32, 3, padding=1)
            self.conv33 = nn.Conv1d(32, 64, 3, padding=1)
            self.pool31 = nn.AvgPool1d(2,2)
            self.pool32 = nn.AvgPool1d(3,3)

            self.flatten = nn.Flatten()
            self.dropout1 = nn.Dropout(0.4)
            self.fc1 = nn.Linear(13312, 4096)
            self.dropout2 = nn.Dropout(0.4)
            self.fc2 = nn.Linear(4096, 1024)
            self.fc3 = nn.Linear(1024, 256)
            self.fc4 = nn.Linear(256, 7)

        def forward(self, x, x0, x1, x2, x3):
            x1 = F.relu(self.conv11(x1))
            x1 = F.relu(self.conv12(x1))
            x1 = self.pool11(x1)
            x1 = F.relu(self.conv13(x1))
            x1 = self.pool12(x1)

            x2 = F.relu(self.conv21(x2))
            x2 = F.relu(self.conv22(x2))
            x2 = self.pool21(x2)
            x2 = F.relu(self.conv23(x2))
            x2 = self.pool22(x2)

            x3 = F.relu(self.conv31(x3))
            x3 = F.relu(self.conv32(x3))
            x3 = self.pool31(x3)
            x3 = F.relu(self.conv33(x3))
            x3 = self.pool32(x3)

            x0 = torch.cat(x1.view(x1.size(0), -1), x2.view(x2.size(0), -1))
            x = torch.cat(x0.view(x0.size(0), -1), x3.view(x3.size(0), -1))
            x = self.dropout1(x)
            x = self.flatten(x)
            x = F.relu(self.fc1(x))
            x = self.dropout2(x)
            x = F.relu(self.fc2(x))
            x = self.fc3(x)
            x = self.fc4(x)

            return x

    return CNN()
